from mindspore import nn
import mindspore.numpy as np


class double_conv(nn.Cell):
    def __init__(self, in_ch, out_ch):
        super(double_conv, self).__init__()
        self.double_conv = nn.SequentialCell(nn.Conv2d(in_ch, out_ch, 3),
                                             nn.BatchNorm2d(out_ch), nn.ReLU(),
                                             nn.Conv2d(out_ch, out_ch, 3),
                                             nn.BatchNorm2d(out_ch), nn.ReLU())

    def construct(self, x):
        x = self.double_conv(x)
        return x


class UNet(nn.Cell):
    def __init__(self, in_ch=1):
        super(UNet, self).__init__()
        # Encoder 编码器
        # [N,3,256,256]->[N,64,255,255]
        # [N, 1, 512, 512] -> [N, 64, 512, 512]
        self.double_conv1 = double_conv(in_ch, 64)
        # [N,64,256,256]->[N,64,128,128]
        # [N,64,512, 512]->[N,64,255, 255]
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # [N,64,128,128]->[N,128,128,128]
        self.double_conv2 = double_conv(64, 128)
        # [N,128,128,128]->[N,128,64,64]
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # [N,128,64,64]->[N,256,64,64]
        self.double_conv3 = double_conv(128, 256)
        # [N,256,64,64]->[N,256,32,32]
        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)

        # [N,256,32,32]->[N,512,32,32]
        self.double_conv4 = double_conv(256, 512)
        # [N,512,32,32]->[N,512,16,16]
        self.maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)
        # [N,512,16,16]->[N,1024,16,16]
        self.double_conv5 = double_conv(512, 1024)

        # Decoder 解码器
        # [N,1024,16,16]->[N,1024,32,32]
        self.upsample1 = nn.ResizeBilinear()
        # [N,1024+512,32,32]->[N,512,32,32]
        self.double_conv6 = double_conv(1024 + 512, 512)
        # [N,512,32,32]->[N,512,64,64]
        self.upsample2 = nn.ResizeBilinear()
        # [N,512+256,64,64]->[N,256,64,64]
        self.double_conv7 = double_conv(512 + 256, 256)
        # [N,256,64,64]->[N,256,128,128]
        self.upsample3 = nn.ResizeBilinear()
        # [N,256+128,128,128]->[N,128,128,128]
        self.double_conv8 = double_conv(256 + 128, 128)
        # [N,128,128,128]->[N,128,256,256]
        self.upsample4 = nn.ResizeBilinear()
        # [N,128+64,256,256]->[N,64,256,256]
        self.double_conv9 = double_conv(128 + 64, 64)

        self.final = nn.Conv2d(64, 3, 1)

    def construct(self, x):
        x1 = self.double_conv1(x)
        tmp = self.maxpool1(x1)

        x2 = self.double_conv2(tmp)
        tmp = self.maxpool2(x2)

        x3 = self.double_conv3(tmp)
        tmp = self.maxpool3(x3)

        x4 = self.double_conv4(tmp)
        tmp = self.maxpool4(x4)

        x5 = self.double_conv5(tmp)

        # nn.ResizeBilinear()
        up_x1 = self.upsample1(x5, scale_factor=2)
        tmp = np.concatenate((x4, up_x1), axis=1)
        tmp = self.double_conv6(tmp)

        up_x2 = self.upsample2(tmp, scale_factor=2)
        tmp = np.concatenate((x3, up_x2), axis=1)
        tmp = self.double_conv7(tmp)

        up_x3 = self.upsample3(tmp, scale_factor=2)
        tmp = np.concatenate((x2, up_x3), axis=1)
        tmp = self.double_conv8(tmp)

        up_x4 = self.upsample4(tmp, scale_factor=2)
        tmp = np.concatenate((x1, up_x4), axis=1)
        tmp = self.double_conv9(tmp)

        # nn.Conv2d(64, 3, 1)
        output = self.final(tmp)
        return output
